import os
import argparse
import torch
import cv2
import numpy as np
from libs.utils import getimages, hairrecolor
import wget
import importlib
import torch.nn.functional as F
os.environ["CUDA_VISIBLE_DEVICES"] = ""  #CPU MODEL ONLY

#modify here
def preprocess(img):
    img = img.astype(np.float32)
    img = img / 255.
    img = (img-0.5)/0.5
    img = np.transpose(img, (2,0,1))
    img = np.expand_dims(img, axis=0)
    return img

def evaluate(net_name, model_path, img_folder, input_shape):
    with torch.no_grad():
        module = importlib.import_module("nets.%s" % net_name)
        model = module.get_segmatting_model()
        model.load_state_dict(torch.load(model_path, map_location=lambda storage, loc: storage))
        model = model.to("cpu")
        model.eval()
        for bgr_img, rgb_img in getimages(img_folder):
            rgb_img = cv2.resize(rgb_img, input_shape)
            processimg = preprocess(rgb_img)
            input_data = torch.from_numpy(processimg).to("cpu")
            input_data = input_data.float()

            mask_pre, alpha_pre = model(input_data)
            scale_mask_pre = F.softmax(mask_pre, dim=1)[0][1]
            scale_mask_pre = scale_mask_pre.detach().numpy()
            alpha_pre = alpha_pre.detach().numpy()[0][0]

            dyeimg = hairrecolor(bgr_img, alpha_pre, (0x40, 0x16, 0x66))

            mask = (255.0*scale_mask_pre).astype(np.uint8)
            alpha = (255*alpha_pre).astype(np.uint8)
            cv2.imshow("img", cv2.resize(bgr_img,(500,500)))
            cv2.imshow("mask", cv2.resize(mask,(500,500)))
            cv2.imshow("alpha", cv2.resize(alpha,(500,500)))
            cv2.imshow("dyeimg", cv2.resize(dyeimg, (500, 500)))
            cv2.waitKey(0)

def main():

    parser = argparse.ArgumentParser()
    parser.add_argument("--net_name", default='ESPNet')
    parser.add_argument("--model_path", default='')
    parser.add_argument("--oss_model_url", default='https://dlmodels.oss-cn-beijing.aliyuncs.com/hairMatting_pytorch/hairnet.pth')
    parser.add_argument("--oss_model_path", default='./sample/models/hairnet.pth')
    parser.add_argument('--img_folder', default='/home/hanson/work/HairMattingTools/train/image')
    parser.add_argument('--input_shape', default=(480, 480) )
    args = parser.parse_args()

    if args.model_path is '':
        if not os.path.exists(args.oss_model_path):
            wget.download(args.oss_model_url, out=args.oss_model_path)
        model_path = args.oss_model_path
    else:
        model_path = args.model_path

    evaluate(args.net_name, model_path, args.img_folder, args.input_shape)


if __name__ == "__main__":
    main()
